import torch
from torch import nn


# 定义一个网络模型
class MyLeNet5(nn.Module):
    # 初始化网络
    def __init__(self):
        super(MyLeNet5, self).__init__()

        # 网络结构：input - c1 -s2 - c3 -s4 -c5- f6- output
        # 只是定义参数方法，所以无需重复sigmoid
        self.c1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=0)
        self.Sigmoid = nn.Sigmoid()
        self.s2 = nn.AvgPool2d(kernel_size=2, stride=2)
        self.c3 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0)
        # self.Sigmoid = nn.Sigmoid()
        self.s4 = nn.AvgPool2d(kernel_size=2, stride=2)
        self.c5 = nn.Conv2d(in_channels=16, out_channels=120, kernel_size=4, stride=1, padding=0)
        # self.Sigmoid = nn.Sigmoid()

        self.flatten = nn.Flatten()  # 平展开
        self.f6 = nn.Linear(in_features=120, out_features=84)
        self.output = nn.Linear(in_features=84, out_features=10)

    # 前向传播
    def forward(self, x):
        x = self.Sigmoid(self.c1(x))
        # print(x.shape)
        x = self.s2(x)
        # print(x.shape)
        x = self.Sigmoid(self.c3(x))
        # print(x.shape)
        x = self.s4(x)
        # print(x.shape)
        x = self.c5(x)
        # print(x.shape)
        x = self.flatten(x)
        # print(x.shape)
        x = self.f6(x)
        # print(x.shape)
        x = self.output(x)
        # print(x.shape)

        return x


if __name__ == "__main__":
    x = torch.randn([1, 1, 28, 28])  # Batch\Channel\Height\Weight
    model = MyLeNet5()  # 实例化网络
    od = model.state_dict()
    print(od.keys())
    model.load_state_dict(torch.load('./save_model/best_model.pth'))
    y = model(x)
    print(y, y.shape)
